package com.hrh.hrhaiagent.agent;

import com.hrh.hrhaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import opennlp.tools.util.StringUtil;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.atomic.AtomicInteger;

@Data
@Slf4j
public abstract class BaseAgent {
    private String name;

    private String systemPrompt;
    private String nextStepPrompt;

    private final AtomicReference<AgentState> state = new AtomicReference<>(AgentState.IDLE);

    private int maxSteps = 10;
    private final AtomicInteger currentStep = new AtomicInteger(0);

    // LLM
    private ChatClient chatClient;

    // Memory
    private final List<Message> messageList = new ArrayList<>();
    private final int maxMessageCount = 100;

    /**
     * 添加消息到上下文，超出最大长度自动移除最早消息
     */
    public synchronized void addMessage(Message msg) {
        if (messageList.size() >= maxMessageCount) {
            messageList.remove(0);
        }
        messageList.add(msg);
    }

    /**
     * 获取当前所有消息的副本
     */
    public synchronized List<Message> getMessages() {
        return new ArrayList<>(messageList);
    }

    /**
     * 清空消息上下文
     */
    public synchronized void clearMessages() {
        messageList.clear();
    }

    /**
     * 运行
     * @param userPrompt
     * @return
     */
    public String run(String userPrompt){
        if (!this.state.compareAndSet(AgentState.IDLE, AgentState.RUNNING)){
            log.warn("[{}] 智能体非空闲，当前状态: {}", name, state.get());
            throw new RuntimeException("Agent is not idle");
        }
        if (StringUtil.isEmpty(userPrompt)){
            this.state.set(AgentState.IDLE);
            log.warn("[{}] 用户输入为空", name);
            throw new RuntimeException("User prompt is empty");
        }
        // 保存上下文
        addMessage(new UserMessage(userPrompt));
        // 保存结果列表
        List<String> results = new ArrayList<>();
        try {
            for (int i = 0; i < maxSteps && state.get() != AgentState.FINISHED; i++){
                int stepNumber = i + 1;
                currentStep.set(stepNumber);
                log.info("[{}] 开始执行第{}步", name, stepNumber);
                String stepResult = step();
                String result = "Step" + stepNumber + ":" + stepResult;
                results.add(result);
            }
            if (currentStep.get() >= maxSteps){
                state.set(AgentState.FINISHED);
                results.add("Agent finished");
            }
            log.info("[{}] 智能体执行完毕，总步数: {}", name, currentStep.get());
            return String.join("\n", results);
        } catch (Exception e) {
            state.set(AgentState.ERROR);
            log.error("[{}] 智能体运行出错，出错步数: {}，错误信息: {}", name, currentStep.get(), e.getMessage(), e);
            return "Agent运行出错: " + e.getMessage();
        } finally {
            this.cleanup();
        }
    }

    /**
     * 运行代理（流式输出）
     *
     * @param userPrompt 用户提示词
     * @return SseEmitter实例
     */
    public SseEmitter runStream(String userPrompt) {
        SseEmitter emitter = new SseEmitter(300000L); // 5分钟超时
        CompletableFuture.runAsync(() -> {
            try {
                if (!this.state.compareAndSet(AgentState.IDLE, AgentState.RUNNING)) {
                    log.warn("[{}] runStream: 智能体非空闲，当前状态: {}", name, state.get());
                    emitter.send("错误：无法从状态运行代理: " + this.state.get());
                    emitter.complete();
                    return;
                }
                if (StringUtil.isEmpty(userPrompt)) {
                    this.state.set(AgentState.IDLE);
                    log.warn("[{}] runStream: 用户输入为空", name);
                    emitter.send("错误：不能使用空提示词运行代理");
                    emitter.complete();
                    return;
                }
                // 记录消息上下文
                addMessage(new UserMessage(userPrompt));
                try {
                    for (int i = 0; i < maxSteps && state.get() != AgentState.FINISHED; i++) {
                        int stepNumber = i + 1;
                        currentStep.set(stepNumber);
                        log.info("[{}] runStream: 执行第{}/{}步", name, stepNumber, maxSteps);
                        String stepResult = step();
                        String result = "Step " + stepNumber + ": " + stepResult;
                        emitter.send(result);
                    }
                    if (currentStep.get() >= maxSteps) {
                        state.set(AgentState.FINISHED);
                        emitter.send("执行结束: 达到最大步骤 (" + maxSteps + ")");
                    }
                    log.info("[{}] runStream: 智能体执行完毕，总步数: {}", name, currentStep.get());
                    emitter.complete();
                } catch (Exception e) {
                    state.set(AgentState.ERROR);
                    log.error("[{}] runStream: 智能体执行出错，出错步数: {}，错误信息: {}", name, currentStep.get(), e.getMessage(), e);
                    try {
                        emitter.send("执行错误: " + e.getMessage());
                        emitter.complete();
                    } catch (Exception ex) {
                        log.error("[{}] runStream: SSE发送错误: {}", name, ex.getMessage(), ex);
                        emitter.completeWithError(ex);
                    }
                } finally {
                    this.cleanup();
                }
            } catch (Exception e) {
                log.error("[{}] runStream: SSE主流程异常: {}", name, e.getMessage(), e);
                emitter.completeWithError(e);
            }
        });
        emitter.onTimeout(() -> {
            this.state.set(AgentState.ERROR);
            this.cleanup();
            log.warn("[{}] SSE连接超时", name);
        });
        emitter.onCompletion(() -> {
            if (this.state.get() == AgentState.RUNNING) {
                this.state.set(AgentState.FINISHED);
            }
            this.cleanup();
            log.info("[{}] SSE连接已完成", name);
        });
        return emitter;
    }

    /**
     * 步骤
     * @return
     */
    public abstract String step();

    /**
     * 清理
     */
    protected void cleanup() {
    }

}
